### 第七章 PyTorch可视化 ###
# 7.1 可视化网络结构
# 7.1.1 使用Print函数打印模型基础信息

import torchvision.models as models
model = models.resnet18()

"""
单纯的print(model)，只能得出基础构件的信息，
既不能显示出每一层的shape，也不能显示对应参数量的大小，
为了解决这些问题，我们就需要介绍出我们今天的主人公torchinfo。
"""

#7.1.2 使用torchinfo可视化网络结构

# torchinfo的安装
# 安装方法一
pip install torchinfo
# 安装方法二
conda install -c conda-forge torchinfo

"""
trochinfo的使用也是十分简单，我们只需要使用torchinfo.summary()就行了，
必需的参数分别是model，input_size[batch_size,channel,h,w]
"""
import torchvision.models as models
from torchinfo import summary
resnet18 = models.resnet18()# 实例化模型
summary(model, (1,3,224,224))  # 1：batch_size 3:图片的通道数 224: 图片的高宽

"""
我们可以看到torchinfo提供了更加详细的信息，包括模块信息（每一层的类型、输出shape和参数量）、
模型整体的参数量、模型大小、一次前向或者反向传播需要的内存大小等

注意：

但你使用的是colab或者jupyter notebook时，
想要实现该方法，summary()一定是该单元（即notebook中的cell）的返回值，
否则我们就需要使用print(summary(...))来可视化。
"""


# 7.2 CNN卷积层可视化

"""

卷积神经网络（CNN）是深度学习中非常重要的模型结构，它广泛地用于图像处理，
极大地提升了模型表现，推动了计算机视觉的发展和进步。
但CNN是一个“黑盒模型”，人们并不知道CNN是如何获得较好表现的，由此带来了深度学习的可解释性问题。
如果能理解CNN工作的方式，人们不仅能够解释所获得的结果，提升模型的鲁棒性，
而且还能有针对性地改进CNN的结构以获得进一步的效果提升。

理解CNN的重要一步是可视化，包括可视化特征是如何提取的、提取到的特征的形式以及模型在输入数据上的关注点等。
本节我们就从上述三个方面出发，介绍如何在PyTorch的框架下完成CNN模型的可视化。

"""

# 7.2.1 CNN卷积核可视化
"""
卷积核在CNN中负责提取特征，可视化卷积核能够帮助人们理解CNN各个层在提取什么样的特征，进而理解模型的工作原理。
例如在Zeiler和Fergus 2013年的paper中就研究了CNN各个层的卷积核的不同，
他们发现靠近输入的层提取的特征是相对简单的结构，而靠近输出的层提取的特征就和图中的实体形状相近了

在PyTorch中可视化卷积核也非常方便，核心在于特定层的卷积核即特定层的模型权重，
可视化卷积核就等价于可视化对应的权重矩阵。
下面给出在PyTorch中可视化卷积核的实现方案，以torchvision自带的VGG11模型为例。
"""
import torch
from torchvision.models import vgg11
model = vgg11(pretrained = True)
print(dict(model.features.named_children()))

# 卷积核对应的应为卷积层（Conv2d），这里以第“3”层为例，可视化对应的参数：
import matplotlib.pyplot as plt #原代码未写

conv1 = dict(model.features.named_children())['3']
kernel_set = conv1.weight.detach()
num = len(conv1.weight.detach())
print(kernel_set.shape)
for i in range(0, num):
    i_kernel = kernel_set[i]
    plt.figure(figsize = (20, 17))
    if (len(i_kernel)) > 1:
        for idx, filter in enumerate(i_kernel):
            plt.subplot(9, 9, idx + 1)
            plt.axis('off')
            plt.imshow(filter[:,:].detach(), cmap = 'bwr')

#torch.Size([128, 64, 3, 3])
# 由于第“3”层的特征图由64维变为128维，因此共有128*64个卷积核

# 7.2.2 CNN特征图可视化方法

"""
与卷积核相对应，输入的原始图像经过每次卷积层得到的数据称为特征图，
可视化卷积核是为了看模型提取哪些特征，可视化特征图则是为了看模型提取到的特征是什么样子的。

获取特征图的方法有很多种，可以从输入开始，逐层做前向传播，直到想要的特征图处将其返回。
尽管这种方法可行，但是有些麻烦了。
在PyTorch中，提供了一个专用的接口使得网络在前向传播过程中能够获取到特征图，这个接口的名称非常形象，叫做hook。
可以想象这样的场景，数据通过网络向前传播，网络某一层我们预先设置了一个钩子，
数据传播过后钩子上会留下数据在这一层的样子，读取钩子的信息就是这一层的特征图。

"""
class Hook(object):
    def __init__(self):
        self.module_name = []
        self.features_in_hook = []
        self.features_out_hook = []

    def __call__(self, module, fea_in, fea_out):
        print("hooker working", self)
        self.module_name.append(module.__class__)
        self.features_in_hook.append(fea_in)
        self.features_out_hook.appemd(fea_out)
        return None

def plot_feature(model, idx, inputs):
    hh = Hook()
    model.feature[idx].register_forward_hook(hh)

    #foward_model(model, False)
    model.eval()
    _ = model(inputs)
    print(hh.module_name)
    print((hh.features_in_hook[0][0].shape))
    print((hh.features_out_hook[0].shape))

    out1 = hh.features_out_hook[0]

    total_ft  = out1.shape[1]
    first_item = out1[0].cpu().clone()

    plt.figure(figsize=(20, 17))


    for ftidx in range(total_ft):
        if ftidx > 99:
            break
        ft = first_item[ftidx]
        plt.subplot(10, 10, ftidx+1)

        plt.axis('off')
        #plt.imshow(ft[ :, :].detach(),cmap='gray')
        plt.imshow(ft[ :, :].detach())

"""
这里我们首先实现了一个hook类，之后在plot_feature函数中，将该hook类的对象注册到要进行可视化的网络的某层中。
model在进行前向传播的时候会调用hook的__call__函数，我们也就是在那里存储了当前层的输入和输出。
这里的features_out_hook 是一个list，每次前向传播一次，都是调用一次，也就是features_out_hook 长度会增加1。
"""

# 7.2.3 CNN class activation map 可视化方法
"""
class activation map （CAM）的作用是判断哪些变量对模型来说是重要的，
在CNN可视化的场景下，即判断图像中哪些像素点对预测结果是重要的。
除了确定重要的像素点，人们也会对重要区域的梯度感兴趣，
因此在CAM的基础上也进一步改进得到了Grad-CAM（以及诸多变种）。
"""
#安装
pip install grad-cam

# 一个简单的例子
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
#以上是必需的
import torch
from torchvision.models import vgg11, resnet18, resnet101, resnext101_32x8d
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np

model = vgg11(pretrained = True)
img_path = "./dog.jpg"
# img_path = 'C:/Users/lenovo/Desktop/dog.jpg'
# resize操作是为了和传入神经网络训练图片大小一致
img = Image.open(img_path).resize((224,224))
# 需要将原始图片转为np.float32格式并且在0-1之间
rgb_img = np.float32(img)/255
plt.imshow(img)

from pytorch_grad_cam import GradCAM,ScoreCAM,GradCAMPlusPlus,AblationCAM,XGradCAM,EigenCAM,FullGrad
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image

target_layers = [model.features[-1]]
# 选取合适的类激活图，但是ScoreCAM和AblationCAM需要batch_size
cam = GradCAM(model=model,target_layers=target_layers)
targets = [ClassifierOutputTarget(preds)]
# 上方preds需要设定，比如ImageNet有1000类，这里可以设为200
grayscale_cam = cam(input_tensor=img_tensor, targets=targets)
# NameError: name 'img_tensor' is not defined
grayscale_cam = grayscale_cam[0, :]
cam_img = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
print(type(cam_img))
Image.fromarray(cam_img)


# 7.2.4 使用FlashTorch快速实现CNN可视化

# 安装
pip install flashtorch -i https://pypi.tuna.tsinghua.edu.cn/simple/ --trusted-host tuna.tsinghua.edu.cn

# 可视化梯度
# Download example images
# !mkdir -p images
# !wget -nv \
#    https://github.com/MisaOgura/flashtorch/raw/master/examples/images/great_grey_owl.jpg \
#    https://github.com/MisaOgura/flashtorch/raw/master/examples/images/peacock.jpg   \
#    https://github.com/MisaOgura/flashtorch/raw/master/examples/images/toucan.jpg    \
#    -P /content/images

import matplotlib.pyplot as plt
import torchvision.models as models
from flashtorch.utils import apply_transforms, load_image
from flashtorch.saliency import Backprop

model = models.alexnet(pretrained = True)
backprop = Backprop(model)

image = load_image('/content/images/great_grey_owl.jpg')
#image = load_image('C:/Users/lenovo/Desktop/dog.jpg')
owl = apply_transforms(image)

target_class = 24
backprop.visualize(owl, target_class, guided = True, use_gpu = True)

"""
运行结果：
D:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py:1025:
UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes
is deprecated and will be removed in future versions. This hook will be missing some grad_input.
Please use register_full_backward_hook to get the documented behavior.
warnings.warn("Using a non-full backward hook when the forward contains multiple autograd Nodes

"""

# 可视化卷积核

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
#以上是必需的

import torchvision.models as models
from flashtorch.activmax import GradientAscent

model = models.vgg16(pretrained=True)
g_ascent = GradientAscent(model.features)

# specify layer and filter info
conv5_1 = model.features[24]
conv5_1_filters = [45, 271, 363, 489]

g_ascent.visualize(conv5_1, conv5_1_filters, title="VGG16: conv5_1")

"""
D:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py:1025:
UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes
is deprecated and will be removed in future versions. This hook will be missing some grad_input.
Please use register_full_backward_hook to get the documented behavior.
warnings.warn("Using a non-full backward hook when the forward contains multiple autograd Nodes "
"""


"""
本节参考
【1】https://andrewhuman.github.io/cnn-hidden-layout_search
【2】https://cloud.tencent.com/developer/article/1747222
【3】https://github.com/jacobgil/pytorch-grad-cam
【4】https://github.com/MisaOgura/flashtorch

"""


# 7.3 使用TensorBoard可视化训练过程


# 7.3.1 TensorBoard安装
 #在已安装PyTorch的环境下使用pip安装即可：

pip install tensorboardX -i https://pypi.tuna.tsinghua.edu.cn/simple/ --trusted-host tuna.tsinghua.edu.cn

# 也可以使用PyTorch自带的tensorboard工具，此时不需要额外安装tensorboard。

# 7.3.2 TensorBoard可视化的基本逻辑

"""
我们可以将TensorBoard看做一个记录员，它可以记录我们指定的数据，
包括模型每一层的feature map，权重，以及训练loss等等。
TensorBoard将记录下来的内容保存在一个用户指定的文件夹里，
程序不断运行中TensorBoard会不断记录。
记录下的内容可以通过网页的形式加以可视化。
"""

# 7.3.3 TensorBoard的配置与启动
"""
在使用TensorBoard前，我们需要先指定一个文件夹供TensorBoard保存记录下来的数据。
然后调用tensorboard中的SummaryWriter作为上述“记录员”
"""

from tensorboardX import SummaryWriter

writer = SummaryWriter('./runs')

"""
上面的操作实例化SummaryWritter为变量writer，并指定writer的输出目录为当前目录下的"runs"目录。
也就是说，之后tensorboard记录下来的内容都会保存在runs。

如果使用PyTorch自带的tensorboard，则采用如下方式import：
"""
from torch.utils.tensorboard import SummaryWriter

"""
是否可以手动往runs文件夹里添加数据用于可视化，或者把runs文件夹里的数据放到其他机器上可视化呢？
答案是可以的。只要数据被记录，你可以将这个数据分享给其他人，其他人在安装了tensorboard的情况下就会看到你分享的数据。

启动tensorboard也很简单，在命令行中输入
"""

tensorboard --logdir=/path/to/logs/ --port=xxxx

"""
其中“path/to/logs/"是指定的保存tensorboard记录结果的文件路径（等价于上面的“./runs"，
port是外部访问TensorBoard的端口号，可以通过访问ip:port访问tensorboard，这一操作和jupyter notebook的使用类似。
如果不是在服务器远程使用的话则不需要配置port。

有时，为了tensorboard能够不断地在后台运行，也可以使用nohup命令或者tmux工具来运行tensorboard。
大家可以自行搜索，这里不展开讨论了。

下面，我们将模拟深度学习模型训练过程，来介绍如何利用TensorBoard可视化其中的各个部分。
"""

# 7.3.4 TensorBoard模型结构可视化
# 首先定义模型：

import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3,out_channels=32,kernel_size = 3)
        self.pool = nn.MaxPool2d(kernel_size = 2,stride = 2)
        self.conv2 = nn.Conv2d(in_channels=32,out_channels=64,kernel_size = 5)
        self.adaptive_pool = nn.AdaptiveMaxPool2d((1,1))
        self.flatten = nn.Flatten()
        self.linear1 = nn.Linear(64,32)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(32,1)
        self.sigmoid = nn.Sigmoid()

    def forward(self,x):
        x = self.conv1(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = self.pool(x)
        x = self.adaptive_pool(x)
        x = self.flatten(x)
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        y = self.sigmoid(x)
        return y

model = Net()
print(model)

"""
输出如下：

Net(
  (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv2): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1))
  (adaptive_pool): AdaptiveMaxPool2d(output_size=(1, 1))
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear1): Linear(in_features=64, out_features=32, bias=True)
  (relu): ReLU()
  (linear2): Linear(in_features=32, out_features=1, bias=True)
  (sigmoid): Sigmoid()
)


可视化模型的思路和7.1中介绍的方法一样，都是给定一个输入数据，前向传播后得到模型的结构，再通过TensorBoard进行可视化，使用add_graph：
"""

writer.add_graph(model, input_to_model = torch.rand(1, 3, 224, 224))
writer.close()

# 展示结果如下（其中框内部分初始会显示为“Net"，需要双击后才会展开）：


# 7.3.5 TensorBoard图像可视化

"""
当我们做图像相关的任务时，可以方便地将所处理的图片在tensorboard中进行可视化展示。

对于单张图片的显示使用add_image
对于多张图片的显示使用add_images
有时需要使用torchvision.utils.make_grid将多张图片拼成一张图片后，用writer.add_image显示
这里我们使用torchvision的CIFAR10数据集为例：
"""

import torchvision
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

transform_train = transforms.Compose(
    [transforms.ToTensor()])
transform_test = transforms.Compose(
    [transforms.ToTensor()])

train_data = datasets.CIFAR10(".", train=True, download=True, transform=transform_train)
test_data = datasets.CIFAR10(".", train=False, download=True, transform=transform_test)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64)

images, labels = next(iter(train_loader))

# 仅查看一张图片
writer = SummaryWriter('./pytorch_tb')
writer.add_image('images[0]', images[0])
writer.close()

# 将多张图片拼接成一张图片，中间用黑色网格分割
# create grid of images
writer = SummaryWriter('./pytorch_tb')
img_grid = torchvision.utils.make_grid(images)
writer.add_image('image_grid', img_grid)
writer.close()

# 将多张图片直接写入
writer = SummaryWriter('./pytorch_tb')
writer.add_images("images",images,global_step = 0)
writer.close()


"""
另外注意上方menu部分，刚刚只有“GRAPHS"栏对应模型的可视化，
现在则多出了”IMAGES“栏对应图像的可视化。左侧的滑动按钮可以调整图像的亮度和对比度。

此外，除了可视化原始图像，TensorBoard提供的可视化方案
自然也适用于我们在Python中用matplotlib等工具绘制的其他图像，用于展示分析结果等内容。
"""

# 7.3.6 TensorBoard连续变量可视化
# TensorBoard可以用来可视化连续变量（或时序变量）的变化过程，通过add_scalar实现：

writer = SummaryWriter('./pytorch_tb')
for i in range(500):
    x = i
    y = x**2
    writer.add_scalar("x", x, i) #日志中记录x在第step i 的值
    writer.add_scalar("y", y, i) #日志中记录y在第step i 的值
writer.close()

"""
如果想在同一张图中显示多个曲线，则需要分别建立存放子路径（使用SummaryWriter指定路径即可自动创建，
但需要在tensorboard运行目录下），同时在add_scalar中修改曲线的标签使其一致即可：
"""

writer1 = SummaryWriter('./pytorch_tb/x')
writer2 = SummaryWriter('./pytorch_tb/y')
for i in range(500):
    x = i
    y = x*2
    writer1.add_scalar("same", x, i) #日志中记录x在第step i 的值
    writer2.add_scalar("same", y, i) #日志中记录y在第step i 的值
writer1.close()
writer2.close()
tb_scalar_twolines

"""
这里也可以用一个writer，但for循环中不断创建SummaryWriter不是一个好选项。
此时左下角的Runs部分出现了勾选项，我们可以选择我们想要可视化的曲线。
曲线名称对应存放子路径的名称（这里是x和y）。

这部分功能非常适合损失函数的可视化，可以帮助我们更加直观地了解模型的训练情况，从而确定最佳的checkpoint。
左侧的Smoothing滑动按钮可以调整曲线的平滑度，当损失函数震荡较大时，将Smoothing调大有助于观察loss的整体变化趋势。

"""


# 7.3.7 TensorBoard参数分布可视化

"""
当我们需要对参数（或向量）的变化，或者对其分布进行研究时，
可以方便地用TensorBoard来进行可视化，通过add_histogram实现。下面给出一个例子：
"""
import torch
import numpy as np

# 创建正态分布的张量模拟参数矩阵
def norm(mean, std):
    t = std * torch.randn((100, 20)) + mean
    return t

writer = SummaryWriter('./pytorch_tb/')
for step, mean in enumerate(range(-10, 10, 1)):
    w = norm(mean, 1)
    writer.add_histogram("w", w, step)
    writer.flush()
writer.close()


# 7.3.8 总结

"""

对于TensorBoard来说，它的功能是很强大的，可以记录的东西不只限于本节所介绍的范围。

主要的实现方案是构建一个SummaryWriter，然后通过add_XXX()函数来实现。

其实TensorBoard的逻辑还是很简单的，它的基本逻辑就是文件的读写逻辑，写入想要可视化的数据，然后TensorBoard自己会读出来。

"""


"""
本节参考
【1】https://blog.csdn.net/Python_Ai_Road/article/details/107704530
"""
